import os
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from flask import current_app
import uuid
from models.unet import UNet
from models.u2net import U2NET
import rasterio
from rasterio.plot import reshape_as_image
from utils.model_downloader import download_model_if_needed

def load_unet_model():
    """
    加载预训练的UNet模型
    
    Returns:
        torch.nn.Module: 加载的模型
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 将输出通道数改为2，以匹配预训练权重
    model = UNet(n_channels=3, n_classes=2)
    
    # 尝试加载模型权重
    model_path = os.path.join(current_app.root_path, 'models', 'weights', 'unet_model.pth')
    
    if not os.path.exists(model_path):
        print(f"UNet模型文件未找到: {model_path}")
        raise FileNotFoundError(f"UNet模型文件未找到，请将模型文件放置到: {model_path}")
    
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()
    return model, device

def load_u2net_model():
    """
    加载预训练的U²-Net模型
    
    Returns:
        torch.nn.Module: 加载的模型
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = U2NET(3, 1)
    
    # 尝试加载模型权重
    model_path = os.path.join(current_app.root_path, 'models', 'weights', 'u2net_model.pth')
    
    if not os.path.exists(model_path):
        print(f"U²-Net模型文件未找到: {model_path}")
        raise FileNotFoundError(f"U²-Net模型文件未找到，请将模型文件放置到: {model_path}")
    
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()
    return model, device

def preprocess_image(image_path, target_size=(512, 512)):
    """
    预处理图像用于模型输入
    
    Args:
        image_path (str): 图像路径
        target_size (tuple): 目标尺寸
    
    Returns:
        torch.Tensor: 预处理后的图像张量
    """
    # 检查文件格式
    if image_path.endswith(('.tif', '.tiff')):
        with rasterio.open(image_path) as src:
            image = src.read()
            image = reshape_as_image(image)
            
            # 如果是多光谱影像，选择RGB波段
            if image.shape[2] > 3:
                image = image[:, :, 0:3]
            
            # 如果不是uint8类型，归一化到0-255
            if image.dtype != np.uint8:
                image = cv2.normalize(image, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    else:
        # 读取常规图像格式
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 调整图像大小
    image = cv2.resize(image, target_size, interpolation=cv2.INTER_AREA)
    
    # 转换为张量并归一化
    image = image / 255.0  # 归一化到0-1
    image = np.transpose(image, (2, 0, 1))  # (H,W,C) -> (C,H,W)
    image = image.astype(np.float32)
    image = torch.from_numpy(image)
    image = image.unsqueeze(0)  # 添加批次维度
    
    return image

def unet_predict(image_path, image_id):
    """
    使用UNet模型进行水域分割预测
    
    Args:
        image_path (str): 图像路径
        image_id (int): 图像ID
    
    Returns:
        tuple: (结果文件名, 水域面积)
    """
    try:
        # 加载模型
        model, device = load_unet_model()
        
        # 预处理图像
        image_tensor = preprocess_image(image_path)
        image_tensor = image_tensor.to(device)
        
        # 进行预测
        with torch.no_grad():
            output = model(image_tensor)
            # 对于2通道输出，我们使用第二个通道（索引1）作为水域的预测
            # 因为通常第一个通道是背景，第二个通道是前景
            if output.shape[1] > 1:  # 多通道输出
                mask = torch.softmax(output, dim=1)[:, 1].squeeze().cpu().numpy()
            else:  # 单通道输出
                mask = torch.sigmoid(output).squeeze().cpu().numpy()
            
            # 二值化
            mask = (mask > 0.5).astype(np.uint8) * 255
        
        # 计算水域面积（像素数）
        water_area = np.sum(mask > 0) / (mask.shape[0] * mask.shape[1])
        
        # 保存结果
        result_filename = f"unet_result_{image_id}_{uuid.uuid4().hex}.png"
        result_path = os.path.join(current_app.config['RESULTS_FOLDER'], result_filename)
        
        # 创建可视化结果
        original_image = cv2.imread(image_path)
        original_image = cv2.resize(original_image, (mask.shape[1], mask.shape[0]))
        
        # 创建叠加效果
        color_mask = np.zeros_like(original_image)
        color_mask[mask > 0] = [0, 0, 255]  # 蓝色表示水域
        
        # 叠加原图和掩码
        alpha = 0.6
        result = cv2.addWeighted(original_image, 1, color_mask, alpha, 0)
        
        # 保存结果
        cv2.imwrite(result_path, result)
        
        return result_filename, float(water_area)
    
    except Exception as e:
        raise RuntimeError(f"UNet预测失败: {str(e)}")

def u2net_predict(image_path, image_id):
    """
    使用U²-Net模型进行水域分割预测
    
    Args:
        image_path (str): 图像路径
        image_id (int): 图像ID
    
    Returns:
        tuple: (结果文件名, 水域面积)
    """
    try:
        # 加载模型
        model, device = load_u2net_model()
        
        # 预处理图像
        image_tensor = preprocess_image(image_path)
        image_tensor = image_tensor.to(device)
        
        # 进行预测
        with torch.no_grad():
            d0, d1, d2, d3, d4, d5, d6 = model(image_tensor)
            # 使用d0作为最终输出（最精细的预测）
            mask = torch.sigmoid(d0).squeeze().cpu().numpy()
            
            # 二值化
            mask = (mask > 0.5).astype(np.uint8) * 255
        
        # 计算水域面积（像素数）
        water_area = np.sum(mask > 0) / (mask.shape[0] * mask.shape[1])
        
        # 保存结果
        result_filename = f"u2net_result_{image_id}_{uuid.uuid4().hex}.png"
        result_path = os.path.join(current_app.config['RESULTS_FOLDER'], result_filename)
        
        # 创建可视化结果
        original_image = cv2.imread(image_path)
        original_image = cv2.resize(original_image, (mask.shape[1], mask.shape[0]))
        
        # 创建叠加效果
        color_mask = np.zeros_like(original_image)
        color_mask[mask > 0] = [0, 0, 255]  # 蓝色表示水域
        
        # 叠加原图和掩码
        alpha = 0.6
        result = cv2.addWeighted(original_image, 1, color_mask, alpha, 0)
        
        # 保存结果
        cv2.imwrite(result_path, result)
        
        return result_filename, float(water_area)
    
    except Exception as e:
        raise RuntimeError(f"U²-Net预测失败: {str(e)}")

def calculate_mndwi(image):
    """
    计算MNDWI (Modified Normalized Difference Water Index)
    
    Args:
        image (numpy.ndarray): 多光谱影像，至少包含绿波段和SWIR波段
        
    Returns:
        numpy.ndarray: MNDWI指数
    """
    # 注意：这里假设绿波段是索引1，SWIR波段是索引5
    # 实际使用时需要根据影像的波段配置进行调整
    if image.shape[2] < 6:
        raise ValueError("影像波段不足，无法计算MNDWI，需要绿波段和SWIR波段")
    
    green = image[:, :, 1].astype(np.float32)
    swir = image[:, :, 5].astype(np.float32)
    
    # 避免除零错误
    denominator = green + swir
    mask = denominator != 0
    
    mndwi = np.zeros_like(green)
    mndwi[mask] = (green[mask] - swir[mask]) / denominator[mask]
    
    return mndwi

def mndwi_calculate(image_path, image_id):
    """
    使用MNDWI方法进行水域提取
    
    Args:
        image_path (str): 图像路径
        image_id (int): 图像ID
    
    Returns:
        tuple: (结果文件名, 水域面积)
    """
    try:
        # 检查文件格式
        if image_path.endswith(('.tif', '.tiff')):
            with rasterio.open(image_path) as src:
                image = src.read()
                image = reshape_as_image(image)
                
                # 如果波段数足够，计算MNDWI
                if image.shape[2] >= 6:
                    # 计算MNDWI
                    mndwi = calculate_mndwi(image)
                    
                    # 二值化，大于0的为水体
                    mask = (mndwi > 0).astype(np.uint8) * 255
                else:
                    # 如果波段不足，使用简化的蓝色通道阈值方法
                    blue_threshold = 150  # 调整阈值
                    if image.shape[2] >= 3:
                        blue = image[:, :, 2]
                        mask = (blue > blue_threshold).astype(np.uint8) * 255
                    else:
                        # 单波段情况
                        mask = (image[:, :, 0] > blue_threshold).astype(np.uint8) * 255
        else:
            # 对于普通RGB图像，使用基于颜色的简单阈值
            img = cv2.imread(image_path)
            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
            
            # 蓝色水域的HSV范围
            lower_blue = np.array([90, 50, 50])
            upper_blue = np.array([130, 255, 255])
            
            # 创建蓝色遮罩
            mask = cv2.inRange(hsv, lower_blue, upper_blue)
        
        # 进行形态学操作，去除噪点
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        
        # 计算水域面积（像素数）
        water_area = np.sum(mask > 0) / (mask.shape[0] * mask.shape[1])
        
        # 保存结果
        result_filename = f"mndwi_result_{image_id}_{uuid.uuid4().hex}.png"
        result_path = os.path.join(current_app.config['RESULTS_FOLDER'], result_filename)
        
        # 创建可视化结果
        original_image = cv2.imread(image_path)
        original_image = cv2.resize(original_image, (mask.shape[1], mask.shape[0]))
        
        # 创建叠加效果
        color_mask = np.zeros_like(original_image)
        color_mask[mask > 0] = [0, 0, 255]  # 蓝色表示水域
        
        # 叠加原图和掩码
        alpha = 0.6
        result = cv2.addWeighted(original_image, 1, color_mask, alpha, 0)
        
        # 保存结果
        cv2.imwrite(result_path, result)
        
        return result_filename, float(water_area)
    
    except Exception as e:
        raise RuntimeError(f"MNDWI计算失败: {str(e)}") 